import numpy as np
import matplotlib.pyplot as plt
from sklearn import svm
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# 用于生成便于二分类的数据集
def generate_data(deviation=1, num_samples_per_class = 100):
    np.random.seed(0)
    mean = [0, 0]  # 均值设为0，对应于原点
    std_deviation = 1.5  # 自定义的标准差
    # 生成第一类样本
    x1 = np.random.normal(loc=mean, scale=deviation, size=(num_samples_per_class, 2)) - [2, 2]
    # 生成第二类样本
    x2 = np.random.normal(loc=mean, scale=deviation, size=(num_samples_per_class, 2)) + [2, 2]
    # 合并样本
    x = np.vstack((x1, x2))
    # 生成标签
    y = np.array([0] * num_samples_per_class + [1] * num_samples_per_class)
    return x, y

# 用于将生成的数据划分成数据集和测试集
def partition_data(X, Y):
    x_train, x_test, y_train, y_test = train_test_split(X, Y, test_size=0.2, random_state=42)
    return x_train, x_test, y_train, y_test

# 用于将模型用matplotlib库绘制出来
def plot_svc_decision_boundary(model, ax=None, plot_support=True):
    if ax is None:
        ax = plt.gca()
    xlim = ax.get_xlim()
    ylim = ax.get_ylim()

    # 创建网格以评估模型
    xx = np.linspace(xlim[0], xlim[1], 30)
    yy = np.linspace(ylim[0], ylim[1], 30)
    YY, XX = np.meshgrid(yy, xx)
    xy = np.vstack([XX.ravel(), YY.ravel()]).T
    P = model.decision_function(xy).reshape(XX.shape)

    # 绘制决策边界和边际
    ax.contour(XX, YY, P, colors="k", levels=[-1, 0, 1], alpha=0.5, linestyles=["--", "-", "--"])

    # 绘制支持向量（如果需要）
    if plot_support:
        ax.scatter(model.support_vectors_[:, 0], model.support_vectors_[:, 1], s=300, linewidth=1, facecolors='none')
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)


if __name__ == "__main__":
    # 步骤1: 生成数据
    X, Y = generate_data()
    # 步骤2: 划分训练集和测试集
    X_train, X_test, Y_train, Y_test = partition_data(X, Y)
    # 步骤3: 训练支持向量机
    clf = svm.SVC(kernel='linear')  # 使用线性核函数
    clf.fit(X_train, Y_train)
    # 步骤4: 绘制数据点及决策边界
    plt.scatter(X[:, 0], X[:, 1], c=Y, cmap=plt.cm.Paired, edgecolors='k')
    plot_svc_decision_boundary(clf)
    plt.title("SVM Linear Classifier")
    plt.show()
    # 步骤5: 测试并评估模型
    Y_pred = clf.predict(X_test)
    print(f"Accuracy: {accuracy_score(Y_test, Y_pred)}")